#!/usr/bin/python

#split dock output into single pdb-files and store additonal information in seperate files

import sys, string, os
from openeye.oechem import *


#Name of output file
output = sys.argv[1]

end_of_molecule = 'true'

output_file = open(output, 'r')

data_list = []
general_list = []
unique_list = []
count_dic = {}

#-------------------------------
def mol2(name):
   
  infile = name
  outfile = name[:-3] + 'mol2'
  
  mol = OEGraphMol()
  
  ifs = oemolistream(infile)
  ofs = oemolostream(outfile)
#  OEReadPDBFile(ifs, mol)
  OEReadMolecule(ifs, mol)
  
  mol.SetTitle(name[:-4])
  
  for atom in mol.GetAtoms():
     res =  OEAtomGetResidue(atom)
     atom.SetPartialCharge(res.GetBFactor())
    
  OEDetermineConnectivity(mol)
  OEPerceiveBondOrders(mol)
  OEFindRingAtomsAndBonds(mol)
  OETriposAtomTypes(mol)
  OETriposAtomTypeNames(mol)
  OETriposBondTypeNames(mol)
  OETriposAtomNames(mol)
  OEAssignAromaticFlags(mol)

  
  OEWriteMol2File(ofs,mol)

#-------------------------------
print "spliting file"


clean_file = []
start = False
end = False
#remove REMfrag fragments
for i in output_file.readlines():
	if i[:6] == 'REMARK':
		start = True
	if start:
		clean_file.append(i)
	if i[:3] == 'TER':
		start = False

#print clean_file

for i in clean_file:
#  print i[0:5]

#if new molecule => create new compound name
  if (end_of_molecule == 'true'):
    compound_name = string.strip(i[67:76])    
    if not compound_name in data_list :
      data_list.append(compound_name)
    else :
      compound_name = compound_name + ' ' +str(data_list.count(acd_label)+1)
      data_list.append(compound_name)
    end_of_molecule = 'false'
#    print compound_name

    acd_label = string.strip(i[67:-1])+'.pdb'
    if not acd_label in unique_list :
      data_list.append(acd_label)
      unique_list.append(acd_label)
    else :
      if not count_dic.has_key(acd_label):
        count_dic[acd_label] = 1
      else :
        count_dic[acd_label] = count_dic[acd_label] + 1
     
      acd_label = acd_label[0:-4] + '_' +str(count_dic[acd_label]) + '.pdb'
      data_list.append(acd_label)
      unique_list.append(acd_label)
    #print acd_label
    molecule_file = open (acd_label, 'w')
    molecule_file.write(i)
    energy = string.strip(i[24:35])
#    print energy
    data_list.append(energy)
    energy_2 = string.strip(i[36:48])
#    print contact
    data_list.append(energy_2)
    energy_3 = string.strip(i[49:58])
#    print nconfs
    data_list.append(energy_3)
    energy_4 = string.strip(i[58:67])
#    print nconfs
    data_list.append(energy_4)

#print coordinates in new file
  elif i[0:4] == 'ATOM':
    #change enlarge space between second last colum and third last colum  => sybyl can read charges
    new_line=i[0:56] + '1.00' +i[56:62]+'\n'
#    new_line=i[0:70]+'\n'
#    print new_line
    molecule_file.write(new_line)

#if molecule ends => close file
  else:
   end_of_molecule = 'true'
   molecule_file.write('TER\n')
   molecule_file.close()
#   print end_of_molecule
   #es soll nach ACD-Nr sortiert werden, da ich nicht weiss, wie man PYTHON
   # sagen kann, dass eine Liste nach dem 2. Wert in einer Unterliste sortiert
   # wird, tausche ich hier einfach den ACD_Label auf Platz 1
   temp = data_list[1]
   data_list[1] = data_list[0]
   data_list[0] = temp
   general_list.append(data_list)
   data_list = []
  
#   convert to mol2 file
   mol2(acd_label)
#   delete pdb file
   command = 'rm -f ' + acd_label
   os.system(command)
 
file_name_file = open ('file_name_list.txt','w')
general_data_file = open ('output_data.txt','w')
general_data_file.write('docking solution\t' + 'ligand\t' + 'elec\t' + 'vdw\t' + 'total\t' + 'rmsd\n')
general_list.sort()
for i in general_list:
  counter = 0
  for j in i:
   counter = counter + 1
#   print j[-4:]
   if j[-4:]== '.pdb':
     file_name_file.write(j[0:-4] + '\n')
     general_data_file.write(j[0:-4] + '\t')
   elif counter < 7:
#     print counter
#     print j
     general_data_file.write(j + '\t') 
   if counter == 6:
     general_data_file.write('\n') 
 
#print general_list
output_file.close()
general_data_file.close()
file_name_file.close()

